import os
import sys
from mmengine.config import Config


def generate_cfg(cfg_file, data_root, classes, save_path):
    '''
    :param cfg_file: default yolov8_s_syncbn_fast_8xb16-500e_coco
    :param data_root:
    :param classes: labels_list
    :param save_path:
    :return: cfg_path
    '''
    cfg = Config.fromfile(cfg_file)

    # About dataroot
    cfg.data_root = data_root
    cfg.train_dataloader.dataset.data_root = cfg.data_root
    cfg.val_dataloader.dataset.data_root = cfg.data_root

    # About data
    # class_name
    cfg.class_name = classes
    # num_classes
    cfg.num_classes = len(cfg.class_name)
    cfg.model.bbox_head.head_module.num_classes = cfg.num_classes
    cfg.model.train_cfg.assigner.num_classes = cfg.num_classes
    # metainfo
    cfg.metainfo = dict(classes=cfg.class_name, palette=[(20, 220, 60),(220, 20, 60)])
    cfg.train_dataloader.dataset.metainfo = cfg.metainfo
    cfg.val_dataloader.dataset.metainfo = cfg.metainfo
    # About ann_file
    cfg.train_ann_file = 'annotations/instances_train2017.json'
    cfg.train_data_prefix = 'train2017/'  # Prefix of train image path
    cfg.val_ann_file = 'annotations/instances_val2017.json'
    cfg.val_data_prefix = 'val2017/'  # Prefix of val image path
    cfg.train_dataloader.dataset.ann_file = cfg.train_ann_file
    cfg.train_dataloader.dataset.data_prefix = dict(img=cfg.train_data_prefix)
    cfg.val_dataloader.dataset.ann_file = cfg.val_ann_file
    cfg.val_dataloader.dataset.data_prefix = dict(img=cfg.val_data_prefix)
    cfg.val_dataloader.dataset.data_root = cfg.data_root
    cfg.val_evaluator.ann_file = cfg.data_root + cfg.val_ann_file

    cfg.test_dataloader = cfg.val_dataloader
    cfg.test_evaluator = cfg.val_evaluator

    # About model
    # model_load
    cfg.load_from = load_from = 'https://download.openmmlab.com/mmyolo/v0/yolov8/yolov8_s_syncbn_fast_8xb16-500e_coco/yolov8_s_syncbn_fast_8xb16-500e_coco_20230117_180101-5aa5f0f1.pth'  # noqa
    # max_epochs
    cfg.max_epochs = 200
    # relevant
    cfg.close_mosaic_epochs = 20
    cfg.val_interval_stage2 = 1
    cfg.default_hooks.param_scheduler.max_epochs = cfg.max_epochs
    cfg.custom_hooks[1].switch_epoch = cfg.max_epochs - cfg.close_mosaic_epochs
    cfg.train_cfg.max_epochs = cfg.max_epochs
    cfg.train_cfg.dynamic_intervals = [((cfg.max_epochs - cfg.close_mosaic_epochs),
                                        cfg.val_interval_stage2)]

    print("generate cfg in", save_path)
    cfg_path = os.path.join(save_path, 'yolov8_s_syncbn_fast_8xb16-500e_coco.py')
    with open(cfg_path, 'w') as f:
        f.write(cfg.pretty_text)

    return cfg_path


def main(args):
    cfg_file = '/home/jacy/PyPjs/detection/mmyolo/configs/yolov8/yolov8_s_syncbn_fast_8xb16-500e_coco.py'

    data_root = f'/home/jacy/GDW/data_{args[1]}/labelme_to_coco' + '/'
    classes = ['neg', 'pos']
    save_path = f'/home/jacy/GDW/data_{args[1]}'
    generate_cfg(cfg_file, data_root, classes, save_path)  # 默认调整为40epochs


if __name__ == '__main__':
    main(sys.argv)
